#!/usr/bin/env python
"""
==============
Weighted Graph
==============
An example using Graph as a weighted network.
"""
# Author: Aric Hagberg (hagberg@lanl.gov)
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import networkx as nx

# G = nx.Graph()
#
# G.add_edge('a', 'b', weight=0.6)
# G.add_edge('a', 'c', weight=0.2)
# G.add_edge('c', 'd', weight=0.1)
# G.add_edge('c', 'e', weight=0.7)
# G.add_edge('c', 'f', weight=0.9)
# G.add_edge('a', 'd', weight=0.3)
#
# elarge = [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] > 0.5]
# esmall = [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] <= 0.5]
#
# pos = nx.spring_layout(G)  # positions for all nodes
#
# # nodes
# nx.draw_networkx_nodes(G, pos, node_size=700)
# # labels
# nx.draw_networkx_labels(G, pos, font_size=20, font_family='sans-serif')
#
# # edges
# nx.draw_networkx_edges(G, pos, edgelist=elarge,
#                        width=6)
#
# nx.draw_networkx_edges(G, pos, edgelist=esmall,
#                        width=6, alpha=0.5, edge_color='b', style='dashed')
#
# plt.axis('off')
# plt.show()



def get_G():
    G = nx.Graph()

    G.add_edge('a', 'b', weight=0.6)
    G.add_edge('a', 'c', weight=0.2)
    G.add_edge('c', 'd', weight=0.1)
    G.add_edge('c', 'e', weight=0.7)
    G.add_edge('c', 'f', weight=0.9)
    G.add_edge('a', 'd', weight=0.3)
    return G

def get_edgelist_by_weight_equal(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] == weight]

def get_edgelist_by_weight_greater(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] > weight]

def get_edgelist_by_weight_greaterequal(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] >= weight]

def get_edgelist_by_weight_less(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] < weight]

def get_edgelist_by_weight_lessequal(G, weight ):
    return [(u, v) for (u, v, d) in G.edges(data=True) if d['weight'] <= weight]

def nx_get_pos(G):
    return nx.spring_layout(G)

def nx_draw_network_nodes_labels(G, pos, _node_size=70, _font_size=20):
    nx.draw_networkx_nodes(G, pos, node_color='g', node_size = _node_size)
    nx.draw_networkx_labels(G, pos, font_color='r', font_size = _font_size, font_family='sans-serif')

def nx_draw_network_edges(G, pos, _edge_list , weight , _edge_color='b',_style='solid'):
    nx.draw_networkx_edges(G, pos, edgelist = _edge_list, width= weight , edge_color=_edge_color, style=_style)

def plt_show():
    plt.axis('off')
    plt.show()


G = get_G()
pos = nx_get_pos(G)
nx_draw_network_nodes_labels(G, pos)
elarge = get_edgelist_by_weight_greater(G, 0.5)
esmall = get_edgelist_by_weight_lessequal(G, 0.5)
nx_draw_network_edges(G, pos, elarge, weight=2, _edge_color='b')
nx_draw_network_edges(G, pos, esmall, weight=1, _edge_color='y', _style='dashed')
# nx_draw_network_edges(G, pos, esmall, weight=1, _edge_color='y', _style='dotted')

# p r  y g b
plt_show()


def plt_networkx_main_first(G):
    pos = nx_get_pos(G)
    nx_draw_network_nodes_labels(G, pos)
    coauthor_edge = get_edgelist_by_weight_equal(G, 3)
    nx_draw_network_edges(G, pos, coauthor_edge, weight=3, _edge_color='b')
    plt_show()

def plt_networkx_main_second(G):
    pos = nx_get_pos(G)
    nx_draw_network_nodes_labels(G, pos)
    coauthor_edge = get_edgelist_by_weight_equal(G, 3)
    second_edge = get_edgelist_by_weight_equal(G, 2)
    nx_draw_network_edges(G, pos, coauthor_edge, weight=3, _edge_color='b')
    nx_draw_network_edges(G, pos, second_edge, weight=2, _edge_color='y', _style='dashed')
    plt_show()

